from flask import Flask, request, jsonify, render_template, Response
import requests
import json
import time

app = Flask(__name__)

def stream_deepseek_response(prompt, history=None):
    """流式响应生成器"""
    url = 'http://localhost:11434/api/generate'
    headers = {'Content-Type': 'application/json'}

    # 构造带历史记录的提示
    full_prompt = f"{history}\nUser: {prompt}" if history else f"User: {prompt}"

    data = {
        "model": "deepseek-r1:14b",
        "prompt": full_prompt,
        "stream": True,  # 启用流式输出,
        "temperature": 0.1,
        "max_tokens": 500,
        "top_p": 0.8,
        "do_sample": False
    }

    try:
        response = requests.post(url, headers=headers, json=data, stream=True)
        response.raise_for_status()

        full_response = []
        for line in response.iter_lines():
            if line:
                chunk = json.loads(line)
                if 'response' in chunk:
                    token = chunk['response']
                    full_response.append(token)
                    yield json.dumps({"token": token, "done": False}) + "\n"

        # 生成最终历史记录
        new_history = f"{full_prompt}\nAssistant: {''.join(full_response)}"
        yield json.dumps({"history": new_history, "done": True}) + "\n"

    except Exception as e:
        yield json.dumps({"error": str(e), "done": True}) + "\n"

@app.route('/')
def index():
    return render_template('indexstream2.html')

@app.route('/query', methods=['POST'])
def handle_query():
    user_input = request.form.get('input')
    history = request.form.get('history', '')

    if user_input.strip().lower() == 'exit':
        return jsonify({"response": "感谢使用DeepSeek对话系统，再见！", "history": ""})

    # 返回流式响应
    return Response(stream_deepseek_response(user_input, history), mimetype='text/event-stream')

if __name__ == "__main__":
    app.run(debug=True)
    # app.run(host='0.0.0.0',port=5001, debug=True)  # 关键修改：允许外部连接
